package util

import (
	"github.com/miekg/dns"
	"net"
	"strings"
)

func ProcessAddress(address string) string {
	parts := strings.Split(address, ":")
	if len(parts) > 2 {
		return address
	}

	host := parts[0]
	port := ""

	if len(parts) == 2 {
		port = parts[1]
	} else {
		port = "80"
	}

	var ip string
	var err error
	if isIP(host) {
		ip = host
	} else {
		if ip, err = LookupIP(host); err != nil {
			panic(err)
		}
	}

	return ip + ":" + port
}

func isIP(ip string) bool {
	if net.ParseIP(ip) == nil {
		return false
	}
	return true
}

// LookupIP looks up host
// Because rocketmq client requires ip address rather than domain address
// It returns a slice of that host's IPv4 addresses.
func LookupIP(host string) (string, error) {
	ips, err := net.LookupIP(host)

	if err != nil {
		// 尝试第二种方式获取ip，适用于k8s环境
		// 设置 DNS 服务器的地址
		dnsServer := "192.168.100.3:53"
		// 执行 DNS 查询
		ipss, err := queryDNS(dnsServer, host)
		if err != nil {
			return "", err
		}
		return ipss[0], nil
	}

	return ips[0].To4().String(), nil
}

// queryDNS 使用 miekg/dns 库进行 DNS 查询
func queryDNS(server, domain string) ([]string, error) {
	client := dns.Client{}
	msg := dns.Msg{}
	msg.SetQuestion(dns.Fqdn(domain), dns.TypeA)

	r, _, err := client.Exchange(&msg, server)
	if err != nil {
		return nil, err
	}

	var ips []string
	for _, ans := range r.Answer {
		if a, ok := ans.(*dns.A); ok {
			ips = append(ips, a.A.String())
		}
	}

	return ips, nil
}
